#!/usr/bin/env python
# -*- coding:utf-8 -*-

"""
    ZhiPu LLM Model
"""

from zhipuai import ZhipuAI


class ZhiPuModel:
    temperature = 0.8
    max_tokens = 8192

    def __init__(self, api_key: str):
        """
            Initialize of the ZhiPu LLM, platform location: https://open.bigmodel.cn/dev
        :param api_key: access key
        """
        self.client = ZhipuAI(api_key=api_key)

    def chat(self,
             messages: list,
             model: str = "glm-4",
             stream: bool = False,
             max_tokens: int = None,
             temperature: float = None) -> str:
        """
            Chat with the ZhiPu LLM
        :param model : model name
        :param messages: list of messages
        :param stream: stream
        :param max_tokens: max tokens
        :param temperature: temperature
        """
        response = self.client.chat.completions.create(
            model=model,
            messages=messages,
            stream=stream,
            max_tokens=self.max_tokens if max_tokens is None else max_tokens,
            temperature=self.temperature if temperature is None else temperature,
        )
        if stream:
            return response

        return response.choices[0].message

    @classmethod
    def get_response(cls, response, stream=True):
        if stream:
            return response.choices[0].delta.content
        else:
            return response.choices[0].message.content
